#!/usr/bin/env python
import sys
import math
import numpy as np
import os, time, random
from multiprocessing import Pool

def process_images(name, task_cmd):
    print('Run task %s (%s)...' % (name, os.getpid()))
    start = time.time()
    print('[RUN CMD]: %s' % (task_cmd))
    os.system(task_cmd)
    end = time.time()
    print('Task %s runs %0.2f seconds.' % (name, (end - start)))

def get_task_infos(src_task_file, origin_t1_dir, processed_t1_dir):
    cmd_list = []

    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")

        input_dir = os.path.join(origin_t1_dir, subj_name)
        iput_file = os.listdir(input_dir)[0]
        full_input_file = os.path.join(input_dir, iput_file)

        recon_cmd = "recon-all -all"
        # -i
        recon_cmd = recon_cmd + " -i " + full_input_file
        # -s
        recon_cmd = recon_cmd + " -s " + subj_name
        # -sd
        recon_cmd = recon_cmd + " -sd " + processed_t1_dir
        cmd_list.append(recon_cmd)

    return cmd_list, len(cmd_list)

if __name__=='__main__':
    print('Parent process %s.' % os.getpid())
    process_num = 11

    project_dir = '/mri_projects/ASD_Analysis'

    src_task_file = os.path.join(project_dir, 'Lists/recon_list_OHSU2.txt')
    origin_t1_dir = os.path.join(project_dir, 'Origin/T1/')
    processed_t1_dir = os.path.join(project_dir, 'DataProcessed/T1/')

    src_lock_file = src_task_file + ".t1.lock."
    if os.path.exists(src_lock_file):
        print("[Notice]: File List `%s` Preporcessing Finished!" % (src_task_file))
    else:
        if not os.path.exists(processed_t1_dir):
            os.makedirs(processed_t1_dir)

        task_list, task_size = get_task_infos(src_task_file, origin_t1_dir, processed_t1_dir)

        p = Pool(process_num)
        for i in range(task_size):
            p.apply_async(process_images, args=(i,task_list[i],))
        print('Waiting for all subprocesses done...')
        p.close()
        p.join()
        print('All subprocesses done.')

        print("########## Lock File ##########")
        print("[Lock T1]: %s" % (src_task_file))
        cp_cmd = "cp " + src_task_file + " " + src_lock_file
        os.system(cp_cmd)
        print("########## Done ##########")
